"""
实验3：记忆系统的内容检索
学生需要使用 LangChain 的 ConversationBufferMemory 管理会话历史
"""
from typing import Dict, List, Any


# 简单的会话记忆类，模拟 ConversationBufferMemory 的行为
class SimpleMemory:
    def __init__(self, memory_key: str = "history"):
        self.memory_key = memory_key
        self.buffer = ""
        self.chat_memory = SimpleChatMemory()
    
    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, Any]):
        # 保存到buffer
        if "input" in inputs and "output" in outputs:
            self.buffer += f"\nHuman: {inputs['input']}\nAI: {outputs['output']}"
            # 保存到chat_memory
            self.chat_memory.add_message("human", inputs['input'])
            self.chat_memory.add_message("ai", outputs['output'])
    
    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        return {self.memory_key: self.buffer}


# 简单的聊天记忆类，模拟 ChatMemory 的行为
class SimpleChatMemory:
    def __init__(self):
        self.messages = []
    
    def add_message(self, role: str, content: str):
        # 创建简单的消息对象
        message = SimpleMessage(role, content)
        self.messages.append(message)


# 简单的消息类
class SimpleMessage:
    def __init__(self, role: str, content: str):
        self.role = role
        self.content = content
        # 为了兼容性，也提供type属性
        self.type = role


# 全局 Memory 映射：{session_id: SimpleMemory 实例}
SESSION_MEMORIES: Dict[str, SimpleMemory] = {}


def chat_with_langchain_memory(message: str, session_id: str) -> dict:
    """
    使用模拟的 ConversationBufferMemory 进行对话
    
    参数:
        message: 用户消息
        session_id: 会话ID
    
    返回:
        字典，包含:
        - response (str): AI回复
        - memory_variables (dict): 会话记忆的内部变量
    
    实现要求:
        1. 使用模拟的 Memory 管理每个 session 的历史
        2. memory_variables 应包含 'history' 键
        3. 不同 session 的 Memory 必须独立
    """
    global SESSION_MEMORIES
    
    # 1. 检查 session_id 是否存在对应的 Memory，不存在则创建
    if session_id not in SESSION_MEMORIES:
        SESSION_MEMORIES[session_id] = SimpleMemory(memory_key="history")
    
    memory = SESSION_MEMORIES[session_id]
    
    # 生成模拟回复
    response_text = f"我已收到您的消息: {message}"
    
    # 保存对话到记忆
    memory.save_context({"input": message}, {"output": response_text})
    
    # 获取memory_variables
    memory_variables = memory.load_memory_variables({})
    
    # 返回响应和 memory_variables
    return {
        "response": response_text,
        "memory_variables": memory_variables
    }


def get_memory_summary(session_id: str) -> str:
    """
    获取指定会话的完整历史记录摘要
    
    参数:
        session_id: 会话ID
    
    返回:
        格式化的历史记录字符串，如:
        "User: 你好\nAI: 你好！有什么可以帮助你的吗？\nUser: ..."
        
        如果会话不存在，返回空字符串或提示信息
    
    实现要求:
        1. 返回人类可读的格式
        2. 包含完整的用户输入和AI回复
        3. 对不存在的 session_id，返回空字符串或提示
    """
    global SESSION_MEMORIES
    
    # 1. 检查 session_id 是否存在
    if session_id not in SESSION_MEMORIES:
        return "会话不存在"
    
    try:
        # 2. 获取 Memory 的 buffer 或 chat_memory
        memory = SESSION_MEMORIES[session_id]
        
        # 获取历史记录
        history = memory.buffer if hasattr(memory, 'buffer') else ""
        
        # 如果buffer不存在，尝试从chat_memory获取
        if not history and hasattr(memory, 'chat_memory'):
            messages = memory.chat_memory.messages
            formatted_messages = []
            for msg in messages:
                if hasattr(msg, 'role') and hasattr(msg, 'content'):
                    role = 'User' if msg.role == 'human' else 'AI' if msg.role == 'assistant' else msg.role
                    formatted_messages.append(f"{role}: {msg.content}")
                elif hasattr(msg, 'type') and hasattr(msg, 'content'):
                    role = 'User' if msg.type == 'human' else 'AI' if msg.type == 'ai' else msg.type
                    formatted_messages.append(f"{role}: {msg.content}")
            history = "\n".join(formatted_messages)
        
        # 3. 格式化消息历史为可读字符串
        # 如果history已经是格式化的字符串，直接返回
        # 否则尝试从memory_variables中获取history
        if not history and hasattr(memory, 'load_memory_variables'):
            memory_vars = memory.load_memory_variables({})
            if 'history' in memory_vars:
                history = memory_vars['history']
        
        # 4. 返回格式化结果
        return history if history else "暂无历史记录"
    except Exception as e:
        return f"获取历史记录失败: {str(e)}"


def clear_memory(session_id: str = None):
    """
    清除会话记忆（辅助函数，用于测试）
    
    参数:
        session_id: 要清除的会话ID，如果为 None 则清除所有会话
    """
    global SESSION_MEMORIES
    
    if session_id is None:
        SESSION_MEMORIES.clear()
    elif session_id in SESSION_MEMORIES:
        del SESSION_MEMORIES[session_id]


# 测试代码（可选，用于学生本地调试）
if __name__ == "__main__":
    # 清空记忆
    clear_memory()
    
    # 测试对话
    print("=== 测试 LangChain Memory ===")
    session_id = "test_session"
    
    result1 = chat_with_langchain_memory("我的电话是13800138000", session_id)
    print(f"第1次对话:")
    print(f"  Response: {result1['response'][:50]}")
    print(f"  Memory variables keys: {result1['memory_variables'].keys()}")
    
    result2 = chat_with_langchain_memory("我的邮箱是test@example.com", session_id)
    print(f"\n第2次对话:")
    print(f"  Response: {result2['response'][:50]}")
    
    # 获取历史摘要
    summary = get_memory_summary(session_id)
    print(f"\n历史摘要:\n{summary}")
    
    # 验证信息是否保存
    print(f"\n验证信息持久化:")
    print(f"  包含电话号码: {'13800138000' in summary}")
    print(f"  包含邮箱: {'test@example.com' in summary}")
